package com.jin.webvideo.config.filter.access.rule;

import com.google.common.cache.Cache;
import com.jin.webvideo.base.util.TokenBucket;
import com.jin.webvideo.config.filter.access.GlobalStore;
import com.jin.webvideo.config.filter.access.handler.AccessHandler;
import com.jin.webvideo.entity.SessionStore;

import javax.servlet.http.HttpServletRequest;

/**
 * 令牌桶规则
 * @author biao532
 * @date 2023/12/01 14:53
 **/
public abstract class TokenBucketRule implements AccessHandler.AccessRule {
    private TokenBucket bucket;
    private long expire;
    private String storeKey;
    public TokenBucketRule(String key, int frequency, int bucketSize) {
        this.bucket = new TokenBucket();
        bucket.setBucketSize(bucketSize);
        bucket.setFrequency(frequency);
        this.storeKey = key;
        this.expire = frequency*bucketSize*2;
    }

    @Override
    public boolean compute(GlobalStore globalStore, SessionStore store, HttpServletRequest request) {
        String key = tokenBucketKey(store, request);
        if (key == null || globalStore == null || storeKey == null) {
            return true;
        }
        Cache<String, TokenBucket> cache = globalStore.getOrCreateExpireCache(storeKey, expire);
        TokenBucket tokenBucket = cache.getIfPresent(key);
        if (tokenBucket == null) {
            tokenBucket = bucket.copy();
            cache.put(key, tokenBucket);
        }
        return tokenBucket.allow();
    }

    abstract String tokenBucketKey(SessionStore store, HttpServletRequest request);
}
